﻿import numpy as np
import random
from collections import defaultdict


class QLearningAgent:
    def __init__(self, env):
        self.env = env
        self.q_table = defaultdict(lambda: np.zeros(env.action_space.n))

    def state_to_key(self, obs):
        """将观察值转换为Q-table的键（简化状态表示）"""
        player_r, player_c = obs["player"]
        # 只关注两个固定奖励点是否存在
        reward_20 = obs["rewards"][2, 0]  # (2,0)位置奖励状态
        reward_22 = obs["rewards"][2, 2]  # (2,2)位置奖励状态
        return (player_r, player_c, reward_20, reward_22)

    def choose_action(self, epsilon=0.1):
        """ε-greedy策略选择动作"""

        obs = self.env._get_obs()
        state_key = self.state_to_key(obs)
        # state_key,
        if random.random() < epsilon:
            return self.env.action_space.sample()  # 随机探索
        else:
            return np.argmax(self.q_table[state_key])  # 选择最优动作

    def update(self, action, reward):
        pass
